import os 


from os.path import join, dirname, abspath,isdir


from tqdm import tqdm



        
    
import json
from visualize import *
import cv2
import numpy as np 

from utils import *
from PIL import Image
import PIL


def concat_images(image_paths, output_path = None, direction='horizontal'):
    """
    Concatenates multiple images and saves the result.

    Parameters:
        image_paths (list): List of image file paths.
        output_path (str): Path to save the concatenated image.
        direction (str): Direction to concatenate ('horizontal' or 'vertical').
    """
    if not image_paths or len(image_paths) < 2:
        raise ValueError("At least two images are required for concatenation.")
    
    # Open images
    
    
    if not isinstance(image_paths[0],PIL.Image.Image):
        images = [Image.open(img) for img in image_paths]
    else:
        images = image_paths



    # alpha = 0.5  
    # images.append(Image.blend(images[-2].convert("RGBA"), images[-1].convert("RGBA"), alpha))
    
    if direction == 'horizontal':
        # Calculate dimensions for the concatenated image
        total_width = sum(img.width for img in images)
        max_height = max(img.height for img in images)
        
        # Create a new blank image with the calculated dimensions
        new_image = Image.new('RGB', (total_width, max_height))
        
        # Paste images side by side
        x_offset = 0
        for img in images:
            new_image.paste(img, (x_offset, 0))
            x_offset += img.width
    
    elif direction == 'vertical':
        # Calculate dimensions for the concatenated image
        max_width = max(img.width for img in images)
        total_height = sum(img.height for img in images)
        
        # Create a new blank image with the calculated dimensions
        new_image = Image.new('RGB', (max_width, total_height))
        
        # Paste images one on top of the other
        y_offset = 0
        for img in images:
            new_image.paste(img, (0, y_offset))
            y_offset += img.height
    
    else:
        raise ValueError("Direction must be 'horizontal' or 'vertical'.")
    
    # Save the concatenated image
    if output_path is not None:
        new_image.save(output_path)
        print(f"Concatenated image saved to {output_path}")
        
    return new_image





class BoosterJsonlExporter:


    def __init__(self, root):
        self.root = root

        self.sample_paths =  [join(root,x)  for x in os.listdir(root) if isdir(join(root,x))]
        


            
        

    def load_sample(self,path):
        cv_file = cv2.FileStorage(join(path,'calib_00-02.xml'), cv2.FILE_STORAGE_READ)    
        f_disp = cv_file.getNode("proj_matL").mat()[0,0]
        cx = cv_file.getNode("proj_matL").mat()[0,2]
        cy = cv_file.getNode("proj_matL").mat()[1,2]
        b_disp = float(cv_file.getNode("baselineLR").real()) / 1000.
        cv_file.release()


        
        # join(path,'mask_00.npy')
        def load_one_side(img_path, depth_path):
            
            imgs = []
            for img_name in [ join(img_path, x) for x in os.listdir(img_path) if x.endswith('.png')]:
                # rgb = cv2.cvtColor(cv2.imread(img_name).astype(np.uint8), cv2.COLOR_BGR2RGB)
                rgb = cv2.imread(img_name).astype(np.uint8)
                imgs.append(rgb)

            turn_to_depth = True
            depth = read_d(depth_path, invert=turn_to_depth, f=f_disp, b=b_disp)
            np.load(depth_path)
            
            mask = np.isinf(depth)
            depth[mask] =0
            
            return imgs,depth, ~mask
        depth_path = join(path,'disp_00.npy')
        img_path = join(path,'camera_00')
        imgs,depth,valid_mask = load_one_side(img_path,depth_path)


        

        #* ignore 
        if False:
                
            join(path,'camera_02')
            join(path,'disp_02.npy')
            # join(path,'mask_02.npy')
            
        

        return imgs,depth,valid_mask

    def transfer_format(self, target_path):

        

        # rgb_path = join(target_path,'images')

        # os.makedirs(hint_path,exist_ok =True)
        # os.makedirs(rgb_path,exist_ok =True)
        

        for sample_path in tqdm(self.sample_paths):
            sample_name = sample_path.split('/')[-1]

            save_path = join(target_path,sample_name)
            os.makedirs(save_path,exist_ok =True)

            imgs,depth,valid_mask = self.load_sample(sample_path)


            for idx, im in enumerate(imgs):
                save_img_path = f"{save_path}/{idx:06d}.png"
                cv2.imwrite(save_img_path, im)

            depth_idx = 0
            
            save_depth_path = f"{save_path}/depth_{depth_idx:06d}.png"
            
            depth = depth * 1000.0
            depth = depth.astype(np.uint16)

            cv2.imwrite(save_depth_path,  depth)
            
    def check_vis(self, target_path):
        os.makedirs(target_path,exist_ok =True)

        # rgb_path = join(target_path,'images')
        # os.makedirs(hint_path,exist_ok =True)
        # os.makedirs(rgb_path,exist_ok =True)
        

        for sample_path in tqdm(self.sample_paths):
            sample_name = sample_path.split('/')[-1]

            

            imgs,depth,valid_mask = self.load_sample(sample_path)

            depth = depth2colordepth(depth )
            imgs +=[depth]

            merged  = concat_images([ Image.fromarray(x) for x in imgs])

            merged.save(join(target_path,f"{sample_name}.jpg"))
            # break


            

    def __call__(self,target_path, save_name= 'train.jsonl'):

        
        ans = []
        for sample_name in os.listdir(target_path):

            if isdir(join(target_path,sample_name)):

                depth_path = join(target_path,sample_name,'depth_000000.png')
                

                images = [x for x in os.listdir(join(target_path,sample_name)) if x.endswith('.png')]

                for img_name in images:

                    if not img_name.startswith('depth_'):
                        
                        ans.append({
                            "conditioning_image": join(target_path,sample_name,img_name),
                            'image':depth_path,
                            'coarse_depth': join(target_path,sample_name,img_name.replace('.png','.npy')),
                            'pred_intrinsic': join(target_path,sample_name,img_name.replace('.png','.json')),
                        })



        #* write jsonl
        
        with open(join(target_path,save_name), 'w') as f:
            for line in ans:
                f.write(json.dumps(line) + '\n')


            

if __name__ == '__main__':




    root = '/share/project/cwm/shaocong.xu/exp/Lotus/data/booster'
    exporter = BoosterJsonlExporter(root)

    target_path = '/share/project/cwm/shaocong.xu/exp/Lotus/data/booster_processed'



    # exporter.transfer_format(target_path)
    exporter(target_path)
    # exporter.check_vis(target_path)



    # imgs,depth,valid_mask  = exporter() #* (valid_mask) == (depth ==0)
    # tmp = depth2colordepth(depth )






    # cv2.imwrite( 'debug.jpg',tmp)
    # Image.fromarray(tmp).save('debug.jpg')
    # cv2.imwrite('debug2.jpg',valid_mask * 255)